public class WeightedQuickUnionUF {
  private int[] id;
  private int[] size; // 各个根节点对应的分量大小
  private int count;

  public WeightedQuickUnionUF(int N) {
    count = N;
    id = new int[N]; // 每个元素的值都指向父节点;默认指向自身节点
    size = new int[N];
    for (int i = 0; i < N; i++) {
      id[i] = i;
      size[i] = 1;
    }
  }

  public int count() {
    return count;
  }

  public boolean connected(int p, int q) {
    return find(p) == find(q);
  }

  public int find(int p) {
    while (p != id[p]) {
      p = id[p];
    }
    return p;
  }

  public void union(int p, int q) {
    int i = find(p);
    int j = find(q);
    if (i == j) {
      return;
    }
    // 将小树的根节点链接到大树的根节点
    if (size[i] < size[j]) {
      id[i] = j;
      size[j] += size[i];
    } else {
      id[j] = i;
      size[i] += size[j];
    }
    count--;
  }
}
